import google.generativeai as genai
from typing import Optional
from ..config import settings


class GeminiManager:
    """Gemini API 管理器"""
    
    def __init__(self):
        genai.configure(api_key=settings.GEMINI_API_KEY)
        self.model = genai.GenerativeModel('gemini-pro')
        self.vision_model = genai.GenerativeModel('gemini-pro-vision')
    
    async def generate_text(self, prompt: str) -> str:
        """生成文本内容"""
        try:
            response = await self.model.generate_content_async(prompt)
            return response.text
        except Exception as e:
            raise Exception(f"文本生成失败: {str(e)}")
    
    async def generate_image(self, prompt: str) -> bytes:
        """生成图片"""
        try:
            response = await self.vision_model.generate_content_async(prompt)
            return response.image
        except Exception as e:
            raise Exception(f"图片生成失败: {str(e)}")
